import os
import shutil
import json
import platform
import yaml
from flask import Flask, request, jsonify
from flask_cors import cross_origin
from ultralytics import YOLO
from result_true import process_annotation_files
import cv2
import numpy as np
from new_compare import app

# app = Flask(__name__)

# 新增的字典表，用于存储图像名称和相关数据
image_map = {}

# 用于存储不同算法标志的结果
algo0_results_storage = []
algo1_results_storage = []

# 模型路径
model_link = 'best_alldata_20240717.pt'
pixel_to_cm = 3.9
sys = platform.system()

# 读取result_dict的内容
with open('./road.yaml', 'r', encoding='utf-8') as f:
    result_dict = yaml.load(f.read(), Loader=yaml.FullLoader)
result_dict = result_dict['names']


def draw_bboxes(image_path, compare_data):
    image = cv2.imread(image_path)
    print("compare_data:", compare_data)  # 打印 compare_data 查看其结构

    # 处理添加的特征
    for feature in compare_data[0]['added_features']:
        bbox = feature['bbox']
        # 解析四个顶点的坐标
        x_coords = [bbox['[x1][y1]'][0], bbox['[x2][y2]'][0], bbox['[x3][y3]'][0], bbox['[x4][y4]'][0]]
        y_coords = [bbox['[x1][y1]'][1], bbox['[x2][y2]'][1], bbox['[x3][y3]'][1], bbox['[x4][y4]'][1]]
        # 计算用于绘制的矩形边界
        xmin, xmax = min(x_coords), max(x_coords)
        ymin, ymax = min(y_coords), max(y_coords)
        # 绘制矩形
        cv2.rectangle(image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 255, 0), 2)  # 绿色表示新增特征

    # 处理移除的特征
    for feature in compare_data[0]['removed_features']:
        bbox = feature['bbox']
        # 解析四个顶点的坐标
        x_coords = [bbox['[x1][y1]'][0], bbox['[x2][y2]'][0], bbox['[x3][y3]'][0], bbox['[x4][y4]'][0]]
        y_coords = [bbox['[x1][y1]'][1], bbox['[x2][y2]'][1], bbox['[x3][y3]'][1], bbox['[x4][y4]'][1]]
        # 计算用于绘制的矩形边界
        xmin, xmax = min(x_coords), max(x_coords)
        ymin, ymax = min(y_coords), max(y_coords)
        # 绘制矩形
        cv2.rectangle(image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255), 2)  # 红色表示移除特征

    return image  # 返回处理后的图像




# 使用示例
# draw_bboxes('path_to_original_image.jpg', compare_data, 'path_to_output_image.jpg')



def calculate_center(bbox):
    x_center = np.mean([bbox['[x1][y1]'][0], bbox['[x2][y2]'][0], bbox['[x3][y3]'][0], bbox['[x4][y4]'][0]])
    y_center = np.mean([bbox['[x1][y1]'][1], bbox['[x2][y2]'][1], bbox['[x3][y3]'][1], bbox['[x4][y4]'][1]])
    return (x_center, y_center)

def calculate_distance_matrix(centers1, centers2):
    if len(centers1) == 0 or len(centers2) == 0:
        return np.array([])  # 返回空数组以避免错误
    centers1 = np.array(centers1)
    centers2 = np.array(centers2)
    return np.linalg.norm(centers2[:, np.newaxis] - centers1, axis=2)
    return list(results.values())



def compare_bboxes(base_data, new_data, base_image_name, new_image_name):
    results = {}

    # 直接使用基准图像和新图像的名称
    image_name1 = base_image_name
    image_name2 = new_image_name

    added_features = []
    removed_features = []

    # 获取所有类别 ID，取并集以便对所有可能的类别进行比较
    base_detection_list = base_data.get(base_image_name, [])
    new_detection_list = new_data.get(new_image_name, [])

    all_class_ids = set([item['class_id'] for item in base_detection_list] + [item['class_id'] for item in new_detection_list])

    # 遍历所有类别 ID
    for class_id in all_class_ids:
        bbox_list1 = [item for item in base_detection_list if item['class_id'] == class_id]
        bbox_list2 = [item for item in new_detection_list if item['class_id'] == class_id]

        # 确保至少有一个非空的列表来获取 class_name
        if bbox_list1:
            class_name = bbox_list1[0]['class_name']
        elif bbox_list2:
            class_name = bbox_list2[0]['class_name']
        else:
            # 如果两个列表都为空，跳过这个类
            continue

        centers1 = [calculate_center(b['bbox']) for b in bbox_list1]
        centers2 = [calculate_center(b['bbox']) for b in bbox_list2]

        # 计算两个集合之间的距离矩阵
        distances = calculate_distance_matrix(centers1, centers2)

        if distances.size == 0:
            # 如果没有匹配，直接根据长度差异来确定是新增还是移除
            if len(centers1) > len(centers2):
                removed_bboxes = bbox_list1
                for bbox in removed_bboxes:
                    removed_features.append({
                        'area': bbox['area'],
                        'bbox': bbox['bbox'],
                        'class_id': class_id,
                        'class_name': class_name
                    })
            elif len(centers2) > len(centers1):
                added_bboxes = bbox_list2
                for bbox in added_bboxes:
                    added_features.append({
                        'area': bbox['area'],
                        'bbox': bbox['bbox'],
                        'class_id': class_id,
                        'class_name': class_name
                    })
        else:
            if len(centers1) > len(centers2):
                # 找到距离最远的几个来标记为移除的特征
                avg_distances_1 = np.mean(distances, axis=1)
                removed_indices = np.argsort(avg_distances_1)[-abs(len(centers1) - len(centers2)):]
                removed_bboxes = [bbox_list1[i] for i in removed_indices]

                for bbox in removed_bboxes:
                    removed_features.append({
                        'area': bbox['area'],
                        'bbox': bbox['bbox'],
                        'class_id': class_id,
                        'class_name': class_name
                    })

            elif len(centers2) > len(centers1):
                # 找到距离最远的几个来标记为新增的特征
                avg_distances_2 = np.mean(distances, axis=0)
                added_indices = np.argsort(avg_distances_2)[-abs(len(centers2) - len(centers1)):]
                added_bboxes = [bbox_list2[i] for i in added_indices]

                for bbox in added_bboxes:
                    added_features.append({
                        'area': bbox['area'],
                        'bbox': bbox['bbox'],
                        'class_id': class_id,
                        'class_name': class_name
                    })

    # 记录比较结果
    results = {
        'image_name1': image_name1,
        'image_name2': image_name2,
        'added_features': added_features,
        'removed_features': removed_features
    }

    # print(results)

    return [results]


def calculate_real_area(x1, y1, x2, y2, x3, y3, x4, y4, pixel_to_cm_ratio):
    pixel_area = abs(
        (x1 * y2 + x2 * y3 + x3 * y4 + x4 * y1) -
        (y1 * x2 + y2 * x3 + y3 * x4 + y4 * x1)
    ) / 2

    cm_area = pixel_area * (pixel_to_cm_ratio ** 2)

    real_area = cm_area / 10000

    return real_area

def predict_img_single(image_path):
    model = YOLO(model_link)
    if sys == "Windows":
        dir_path = '\\'.join(image_path.split('\\')[:-1])
        image_name = image_path.split('\\')[-1]
    else:
        dir_path = '/'.join(image_path.split('/')[:-1])
        image_name = image_path.split('/')[-1]
    results = model.predict(device="cpu", source=image_path,   save_txt=True,save=True,  conf=0.5, show_labels=False, iou=0.5)
    save_dir = results[0].save_dir

    print(f"save_dir: {save_dir}")

    image_folder = save_dir
    annotation_folder = os.path.join(save_dir, 'labels')
    process_annotation_files(image_folder, annotation_folder)
    result_txt = os.path.join(save_dir, 'labels', f"{image_name.split('.')[0]}_transformed.txt")
    image_data = []
    try:
        with open(result_txt) as fp:
            predict_result = fp.readlines()
        current_path = os.path.dirname(__file__)
        image_link = os.path.abspath(os.path.join(current_path, image_name))
        predict_result = [result.replace('\n', '').split(' ') for result in predict_result]
    except:
        predict_result = []
    for result in predict_result:
        bbox = {
            "[x1][y1]": [float(result[1]), float(result[2])],
            "[x2][y2]": [float(result[3]), float(result[4])],
            "[x3][y3]": [float(result[5]), float(result[6])],
            "[x4][y4]": [float(result[7]), float(result[8])],
        }
        area = calculate_real_area(float(result[1]), float(result[2]), float(result[3]), float(result[4]),
                                   float(result[5]), float(result[6]), float(result[7]), float(result[8]), pixel_to_cm)
        image_data.append({
            "class_id": int(result[0]),
            "class_name": result_dict[int(result[0])],
            "bbox": bbox,
            "area": area
        })

    data = {
        image_name: image_data
    }
    return data, save_dir



@app.route('/upload', methods=['POST'])
@cross_origin(supports_credentials=True)
def upload():
    try:
        json_file_path = request.json.get("json_file_path")

        if not json_file_path or not os.path.exists(json_file_path):
            return jsonify({"code": 400, "message": "Invalid or missing JSON file path"}), 400

        with open(json_file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)

        if 'results' not in data:
            return jsonify({"code": 400, "message": "Invalid JSON format"}), 400

        for result in data['results']:
            file_path = result.get('file_path')
            correct_results = result.get('correct_results', [])
            for item in correct_results:
                image_name = item.get('image_name')
                datas = item.get('datas')
                if image_name and datas:
                    if image_name not in image_map:
                        image_map[image_name] = []
                    image_map[image_name].append(datas)

        print("Image Map:", image_map)


        return jsonify({"code": 200, "message": "Data processed successfully", "image_map": image_map}), 200

    except Exception as e:
        return jsonify({"code": 500, "message": str(e)}), 500

@app.route('/process_image', methods=['POST'])
@cross_origin(supports_credentials=True)
def process_image():
    try:
        image_file_path = request.json.get("image_file_path")
        out_file_path = request.json.get("out_file_path")
        algo_flag = request.json.get("algo_flag")
        compared_file_path = request.json.get("compared_file_path", None)  # 新参数

        if not image_file_path or not out_file_path or algo_flag is None:
            return jsonify({"code": 400, "message": "Missing required parameters"}), 400

        if not os.path.exists(image_file_path):
            return jsonify({"code": 400, "message": "Image file path does not exist"}), 400

        if algo_flag == 0:
            detection_data, save_dir = predict_img_single(image_file_path)
            print(f"Detection save_dir: {save_dir}")
            processed_image_path = os.path.join(save_dir, os.path.basename(image_file_path))
            if os.path.exists(processed_image_path):
                shutil.move(processed_image_path, out_file_path)
            else:
                return jsonify({"code": 500, "message": "Processed image not found"}), 500
            message = "Detection task completed successfully"
            data = {"image_map": detection_data}

            # 格式化结果为需要的格式
            algo0_results_storage.append({
                "file_path": image_file_path,
                "correct_results": [
                    {
                        "image_name": os.path.basename(image_file_path),
                        "datas": detection_data[os.path.basename(image_file_path)]
                    }
                ]
            })
            return jsonify({"code": 200, "message": message, "output_path": out_file_path, "data": data}), 200

        elif algo_flag == 1:
            # 获取时间参数
            base_task_total_time = request.json.get("base_task_total_time")
            new_task_start_time = request.json.get("new_task_start_time")
            change_detection_time = request.json.get("change_detection_time")
            compared_file_path = request.json.get("compared_file_path", None)  # 确保获取 compared_file_path

            # 验证参数是否存在
            if not all([base_task_total_time, new_task_start_time, change_detection_time]):
                return jsonify({"code": 400, "message": "Missing required time parameters for change detection"}), 400

            # 将时间参数转换为浮点数
            try:
                base_task_total_time = float(base_task_total_time)
                new_task_start_time = float(new_task_start_time)
                change_detection_time = float(change_detection_time)
            except ValueError:
                return jsonify({"code": 400, "message": "Time parameters must be numeric"}), 400

            # 计算新任务中已用的时间
            elapsed_time = change_detection_time - new_task_start_time

            # 检查文件是否存在，如果不存在则创建文件
            if not os.path.exists(compared_file_path):
                os.makedirs(os.path.dirname(compared_file_path), exist_ok=True)
                with open(compared_file_path, 'w') as fp:
                    pass
                # 对 image_map 的键（图片名称）进行排序
            ordered_image_names = sorted(image_map.keys())
            print(f"Length of ordered_image_names: {len(ordered_image_names)}")

            # 基准图像的数量
            num_base_images = len(ordered_image_names)

            if num_base_images == 0:
                return jsonify({"code": 400, "message": "No base images available for comparison"}), 400

            # 计算每张基准图像对应的时间间隔
            interval = base_task_total_time / num_base_images
            if interval == 0:
                return jsonify({"code": 400,"message": "Invalid interval calculation, base_task_total_time might be too small"}), 400

            # 计算基准图像的索引
            index = int(elapsed_time / interval)

            # 确保索引在有效范围内
            if index >= num_base_images:
                index = num_base_images - 1
            elif index < 0:
                index = 0
            print(f"Adjusted index: {index}")

            if index >= len(ordered_image_names):
                return jsonify({"code": 500, "message": "Calculated index is out of range"}), 500



            # 获取基准图像名称
            base_image_name = ordered_image_names[index]

            # 检查基准图像是否存在于 image_map 中且数据有效
            if base_image_name not in image_map:
                return jsonify(
                    {"code": 500, "message": f"Base image name '{base_image_name}' not found in image_map"}), 500





            # 从 image_map 中获取基准图像的检测数据
            base_detection_data = {base_image_name: image_map[base_image_name]}
            # 打印用于调试的信息
            print(f"Base image name: {base_image_name}")
            print(f"Detection data for base image: {base_detection_data}")

            # 对新图像进行检测
            new_detection_data, save_dir = predict_img_single(image_file_path)

            # 移动处理后的新图像
            processed_image_path = os.path.join(save_dir, os.path.basename(image_file_path))
            if os.path.exists(processed_image_path):
                shutil.move(processed_image_path, out_file_path)
            else:
                return jsonify({"code": 500, "message": "Processed image not found"}), 500

            # 构建新图像的检测数据字典
            new_image_name = os.path.basename(image_file_path)
            new_detection_data_dict = {new_image_name: new_detection_data[new_image_name]}
            print(f"Detection data for new image: {new_detection_data_dict}")

            # 执行变化检测
            compare_data = compare_bboxes(base_detection_data, new_detection_data_dict,base_image_name,new_image_name)

            # 绘制变化检测结果并保存
            processed_image = draw_bboxes(image_file_path, compare_data)
            cv2.imwrite(compared_file_path, processed_image)

            # 存储结果
            algo1_results_storage.extend(compare_data)
            message = "Change detection task completed successfully"
            data = {"image_map": new_detection_data_dict, "compare_data": compare_data}

            return jsonify({"code": 200, "message": message, "output_path": out_file_path, "data": data}), 200
        else:
            return jsonify({"code": 400, "message": "Invalid algo_flag value"}), 400

    except Exception as e:
        return jsonify({"code": 500, "message": str(e)}), 500


@app.route('/process_json', methods=['POST'])
@cross_origin(supports_credentials=True)
def process_json():
    try:
        algo0_json_path = request.json.get("algo0_json_path")
        algo1_json_path = request.json.get("algo1_json_path")

        if not algo0_json_path or not algo1_json_path:
            return jsonify({"code": 400, "message": "Missing JSON file paths"}), 400

        if not algo0_results_storage and not algo1_results_storage:
            return jsonify({"code": 400, "message": "No results available"}), 400

        # 格式化 algo_flag == 0 的结果
        if algo0_results_storage:
            formatted_algo0_results = {
                "results": []
            }
            for result in algo0_results_storage:
                formatted_algo0_results["results"].append(result)

            with open(algo0_json_path, 'w', encoding='utf-8') as f:
                json.dump(formatted_algo0_results, f, ensure_ascii=False, indent=4)
            algo0_results_storage.clear()  # 清空存储

        # 写入 algo_flag == 1 的结果
        if algo1_results_storage:
            with open(algo1_json_path, 'w', encoding='utf-8') as f:
                json.dump({"results": algo1_results_storage}, f, ensure_ascii=False, indent=4)
            algo1_results_storage.clear()  # 清空存储

        return jsonify({"code": 200, "message": "JSON files created successfully"}), 200

    except Exception as e:
        return jsonify({"code": 500, "message": str(e)}), 500


if __name__ == '__main__':
    from gevent import pywsgi

    server = pywsgi.WSGIServer(('0.0.0.0', 6000), app)
    print('服务器已创建，正在监听端口 6000...', flush=True)
    server.serve_forever()
